from typing import List, Dict
import datetime
import json
import time
import traceback

from static.cookie_pool import cookies

from qdata.baidu_index.common import check_keywords_exists, get_encrypt_json
from qdata.errors import QdataError, ErrorCode

# ALL_KIND = ['all', 'pc', 'wise']
ALL_KIND = ['all']


def get_search_index(
    *,
    keywords_list: List[List[str]],
    start_date: str,
    end_date: str,
    cookies: str,
    area: int = 0
):
    if len(keywords_list) > 5:
        raise QdataError(ErrorCode.KEYWORD_LIMITED)
    start_date = datetime.datetime.strptime(start_date, '%Y-%m-%d')
    end_date = datetime.datetime.strptime(end_date, '%Y-%m-%d')
    encrypt_json = get_encrypt_json(
        start_date=start_date,
        end_date=end_date,
        keywords=keywords_list,
        type='search',
        area=area,
        cookies=cookies
    )
    encrypt_datas = encrypt_json['data']['generalRatio']

    for encrypt_data in encrypt_datas:
        encrypt_data['year'] = start_date.year
        for formatted_data in format_data(encrypt_data):
            yield formatted_data



def format_data(data: Dict):
    """
        格式化堆在一起的数据
    """
    keyword = str(data['word'])


    index_all_data = data['all']['avg']
    index_pc_data = data['pc']['avg']
    index_wise_data = data['wise']['avg']
    formatted_data = {
        'keyword': [keyword_info['name'] for keyword_info in json.loads(keyword.replace('\'', '"'))],
        'all': index_all_data if index_all_data else 0,
        'pc': index_pc_data if index_pc_data else 0,
        'wise': index_wise_data if index_wise_data else 0,
    }

    yield formatted_data

def decrypt_func(kw, kind, encrypt_data):
    """
        选择kw为word的数据的三条属性
    """
    kw_data = [d for d in encrypt_data if 'word' in d and d['word'][0]['name'] == kw]

    if kind in kw_data[0]:
        return kw_data[0][kind]
    else:
        raise QdataError(ErrorCode.KIND_ERROR)
    
    
def get_clear_keywords_list(keywords_path, output_path=None) -> List[List[str]]:
    """
    读取 keywords 文件，检查关键词是否存在，返回有效的关键词列表并保存
    
    Args:
        keywords_path: 输入关键词文件路径
        output_path: 输出清理后关键词的文件路径，如果为None则保存到原文件名_cleared.txt
    
    Returns:
        有效的关键词列表
    """
    # 读取 keywords 文件
    with open(keywords_path, 'r', encoding='utf-8') as f:
        content = f.read().strip()
    
    # 解析关键词，按逗号分割
    if not content:
        return []
    
    keywords = [kw.strip() for kw in content.split(',') if kw.strip()]
    if not keywords:
        return []
    
    # 将关键词分组，每组一个关键词（可以根据需要调整分组策略）
    keywords_list = [[keyword] for keyword in keywords]
    
    # 检查所有关键词是否存在
    exists_keywords = []
    not_exists_keywords = []
    
    # 如果关键词超过15个，需要分批处理
    if len(keywords) <= 15:
        try:
            check_result = check_keywords_exists(keywords, cookies)
            exists_keywords.extend(check_result["exists_keywords"])
            not_exists_keywords.extend(check_result["not_exists_keywords"])
        except Exception:
            traceback.print_exc()
    else:
        # 分批处理，每批最多15个关键词
        for start in range(0, len(keywords), 15):
            batch_keywords = keywords[start:start + 15]
            try:
                check_result = check_keywords_exists(batch_keywords, cookies)
                exists_keywords.extend(check_result["exists_keywords"])
                not_exists_keywords.extend(check_result["not_exists_keywords"])
                time.sleep(5)  # 避免请求过于频繁
            except Exception:
                traceback.print_exc()
                time.sleep(30)  # 出错时等待更长时间

    # 使用已经验证存在的关键词
    
    # 保存清理后的关键词
    if output_path is None:
        # 如果没有指定输出路径，则在原文件名基础上添加_cleared后缀
        import os
        base_name = os.path.splitext(keywords_path)[0]
        ext = os.path.splitext(keywords_path)[1]
        output_path = f"{base_name}_cleared{ext}"
    
    # 将有效关键词保存到文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(','.join(exists_keywords))
    
    print(f"原始关键词数量: {len(keywords)}")
    print(f"有效关键词数量: {len(exists_keywords)}")
    print(f"被过滤的关键词数量: {len(not_exists_keywords)}")
    print(f"清理后的关键词已保存到: {output_path}")
    
    if not_exists_keywords:
        print(f"被过滤的关键词: {', '.join(sorted(not_exists_keywords))}")

    return exists_keywords

if __name__ == "__main__":
    get_clear_keywords_list("static/keywords.txt", "static/keywords_cleared.txt")
